#include <ros/ros.h>
#include <geometry_msgs/Twist.h>
#include <moveit/move_group_interface/move_group_interface.h>
#include <tf2_eigen/tf2_eigen.h>
#include <eigen_conversions/eigen_msg.h>
#include <Eigen/Dense>

class AdmittanceController {
public:
    AdmittanceController() 
        : move_group("arm"), loop_rate(100) {
        move_group.setPlanningTime(10.0);

        // 初始化参数
        M_a_ = Eigen::Matrix<double, 6, 6>::Identity();
        K_a_ = 100 * Eigen::Matrix<double, 6, 6>::Identity();
        D_a_ = 10 * Eigen::Matrix<double, 6, 6>::Identity();
        arm_desired_twist_final_.setZero();
    }

    void computeAdmittance() {
        // 获取机械臂的实际状态
        getArmState();
        
        // 定义期望的加速度向量
        Eigen::Matrix<double, 6, 1> arm_desired_acceleration;
        
        // 定义误差向量
        Eigen::Matrix<double, 6, 1> error;
        
        // 平衡位置和实际位置的之间的误差
        Eigen::Quaterniond quat_rot_err(arm_real_orientation_ * equilibrium_orientation_.inverse());
        Eigen::AngleAxisd err_arm_des_orient(quat_rot_err);
        error.bottomRows(3) << err_arm_des_orient.axis() * err_arm_des_orient.angle();
        
        // 平移误差
        error.topRows(3) = arm_real_position_ - equilibrium_position_;
        
        // 计算期望的加速度
        arm_desired_acceleration = M_a_.inverse() * ( - K_a_ * error - D_a_ * arm_desired_twist_final_ + wrench_external_);
        
        // 计算期望的速度
        arm_desired_twist_final_ += arm_desired_acceleration * loop_rate.expectedCycleTime().toSec();

        // 使用MoveIt控制机械臂
        moveArm();
    }

private:
    void getArmState() {
        // 获取机械臂的实际位置和姿态
        geometry_msgs::PoseStamped current_pose = move_group.getCurrentPose();
        tf2::fromMsg(current_pose.pose.position, arm_real_position_);
        tf2::fromMsg(current_pose.pose.orientation, arm_real_orientation_);
    }

    void getExternalWrench() {
        // 获取外部作用的力和力矩（这里假设为零）
        wrench_external_.setZero();
    }

    void moveArm() {
        geometry_msgs::Twist arm_cmd;
        arm_cmd.linear.x = arm_desired_twist_final_(0);
        arm_cmd.linear.y = arm_desired_twist_final_(1);
        arm_cmd.linear.z = arm_desired_twist_final_(2);
        arm_cmd.angular.x = arm_desired_twist_final_(3);
        arm_cmd.angular.y = arm_desired_twist_final_(4);
        arm_cmd.angular.z = arm_desired_twist_final_(5);

        move_group.setMaxVelocityScalingFactor(1.0);
        move_group.setMaxAccelerationScalingFactor(1.0);

        moveit::core::RobotStatePtr current_state = move_group.getCurrentState();
        const robot_state::JointModelGroup* joint_model_group = current_state->getJointModelGroup(move_group.getName());
        
        current_state->copyJointGroupPositions(joint_model_group, joint_group_positions);
        // 计算新的关节位置
        for (size_t i = 0; i < joint_group_positions.size(); ++i) {
            joint_group_positions[i] += arm_desired_twist_final_(i) * loop_rate.expectedCycleTime().toSec();
        }

        move_group.setJointValueTarget(joint_group_positions);
        move_group.move();
    }

    // 定义机械臂控制相关变量
    moveit::planning_interface::MoveGroupInterface move_group;
    ros::Rate loop_rate;
    Eigen::Vector3d arm_real_position_;
    Eigen::Quaterniond arm_real_orientation_;
    Eigen::Vector3d equilibrium_position_;
    Eigen::Quaterniond equilibrium_orientation_;
    Eigen::Matrix<double, 6, 1> wrench_external_;
    Eigen::Matrix<double, 6, 1> arm_desired_twist_final_;
    Eigen::Matrix<double, 6, 6> M_a_, K_a_, D_a_;
    std::vector<double> joint_group_positions;
};

int main(int argc, char **argv) {
    ros::init(argc, argv, "admittance_controller_node");
    ros::AsyncSpinner spinner(1);
    spinner.start();

    AdmittanceController admittance_controller;

    ros::Rate loop_rate(100); // 100 Hz的循环频率

    while (ros::ok()) {
        admittance_controller.computeAdmittance();
        ros::spinOnce();
        loop_rate.sleep();
    }

    return 0;
}
